import copy


from init import *
import pandas as pd

plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']
plt.rcParams['axes.unicode_minus'] = False


def secret_msg(row, col):
    '''

    :param ec: 需要嵌入的容量
    :return: 生成随机生成的二进制
    '''
    arr = np.arange(row * col)
    np.random.shuffle(arr)
    secret_msg = np.zeros((row, col))
    for i in range(row):
        for j in range(col):
            # secret_msg[i, j] = np.mod(arr[col*(i+1)-j-1],2)
            secret_msg[i, j] = np.mod(arr[col * (i + 1) - j - 1], 2)
    return secret_msg


# 插入修改
def inser_msg(em, a, b, img_hide, predict_value, secrt_msg, flag):
    '''
    :param hs: 误差矩阵e
    :param a: 嵌入左区间
    :param b: 嵌入又区间
    :param secrt_msg: 嵌入信息
    :param flag 判断是插入为奇数还是偶数
    :return:
    '''
    hs = copy.deepcopy(em)
    row, col = hs.shape
    index = 0
    test = 0
    for i in range(1, row - 2):
        for j in range(1, col - 2):
            if np.mod(i + j, 2) == 0 and flag == 0:
                if hs[i][j] == a:
                    hs[i][j] -= secrt_msg[0][index]
                    test += 1
                else:
                    if hs[i][j] == b:
                        hs[i][j] += secrt_msg[0][index]
                    else:
                        if hs[i][j] < a:
                            hs[i][j] -= 1
                        else:
                            if hs[i][j] > b:
                                hs[i][j] += 1
                img_hide[i][j] = predict_value[i][j] + hs[i][j]
            if np.mod(i + j, 2) == 1 and flag == 1:
                if hs[i][j] == a:
                    hs[i][j] -= secrt_msg[0][index]
                    test += 1
                else:
                    if hs[i][j] == b:
                        hs[i][j] += secrt_msg[0][index]
                    else:
                        if hs[i][j] < a:
                            hs[i][j] -= 1
                        else:
                            if hs[i][j] > b:
                                hs[i][j] += 1
                img_hide[i][j] = predict_value[i][j] + hs[i][j]
            index += 1
    cv2.imwrite('./output/out.bmp', img_hide)
    return hs


def insert_msg(em, Dbis, bin, img_hide3, secret, ec3, flag3):
    '''
        em误差矩阵
        Dbis复杂度矩阵
        bin最佳区间
        img_hide图形矩阵
        secret嵌入矩阵
        插入需要插入的信息
    '''
    index = 0
    row, col = img_hide3.shape
    for i in range(1, row - 2):
        for j in range(1, col - 2):
            if np.mod(i + j, 2) == 0 and flag3 == 0 and ec3 > 0:
                index += 1
                ks = Dbis[i][j] - 1
                if em[i][j] == bin[int(ks)]:
                    ec3 = ec3 - 1
                    img_hide3[i][j] += secret[i][j]
                if em[i][j] == -bin[int(ks)] - 1:
                    ec3 = ec3 - 1
                    img_hide3[i][j] -= secret[i][j]
                if em[i][j] > bin[int(ks)]:
                    img_hide3[i][j] += 1
                if em[i][j] < -bin[int(ks)] - 1:
                    img_hide3[i][j] -= 1
            if np.mod(i + j, 2) == 1 and flag3 == 1 and ec3 > 0:
                index += 1
                ks = Dbis[i][j] - 1
                if em[i][j] == bin[int(ks)]:
                    ec3 = ec3 - 1
                    img_hide3[i][j] += secret[i][j]
                if em[i][j] == -bin[int(ks)] - 1:
                    ec3 = ec3 - 1
                    img_hide3[i][j] -= secret[i][j]
                if em[i][j] > bin[int(ks)]:
                    img_hide3[i][j] += 1
                if em[i][j] < -bin[int(ks)] - 1:
                    img_hide3[i][j] -= 1
    return img_hide3, ec3


def Lsb_insert(row, col, anx, img_hide2, flags, anx_secret):
    for i in range(1, row - 2):
        for j in range(1, col - 2):
            if np.mod(i + j, 2) == 0 and flags == 0:
                if anx > 0:
                    img_hide2[i][j] = 2 * np.floor(img_hide2[i][j] / 2) + anx_secret[i][j]
                    anx = anx - 1
            if np.mod(i + j, 2) == 1 and flags == 1:
                if anx > 0:
                    img_hide2[i][j] = 2 * np.floor(img_hide2[i][j] / 2) + anx_secret[i][j]
                    anx = anx - 1
    return img_hide2


def danci(danci_img, flags, k, ecfinals):
    danci_img, lcm, halfsize, data = Over_flow(danci_img, flags)
    # 数组长度
    data_length = len(data)
    # 获取图片的坐标的长宽
    row, col = danci_img.shape
    # 初始图像数据,获取误差以及按照图像的复杂度进行分类
    Complexity, em, predict_value = img_division(danci_img, flags)
    # 计算生成二维多重直方图
    hist = divide_hs(Complexity, em, flags)
    pd.DataFrame(em).to_csv('./em1.csv')
    pd.DataFrame(hist).to_csv('./hist1.csv')
    pd.DataFrame(Complexity).to_csv('./complex1.csv')
    # print('hist[1:2,1]',hist[0,0:255])
    # 归一化
    T = guiyi(16, Complexity, hist, halfsize)
    # 重新计算负载度
    Dbis = recomplex(row, col, 16, T, flags, Complexity)
    pd.DataFrame(Dbis).to_csv('./Dbis.csv')
    # 新的二重直方图
    new_hist = re_hs(16, row, col, Dbis, flags, em)
    pd.DataFrame(new_hist).to_csv('./newhist.csv')
    # 辅助信息
    print("t", T)
    s1 = np.ceil(np.log2(T[k - 2]))
    k = 16
    anx = cal_aux1(s1, k, lcm, T, data_length)
    print(anx)
    # 当前嵌入容量
    ec1 = ecfinals / 2 + anx
    print('anx,', anx)
    print(ec1)
    # 找到最佳嵌入点
    bins = searchbins(new_hist, ec1)
    print('bin', bins)
    '''
        通过最佳嵌入点判断是否存储了两个以上的250（即无穷大）
        如果超过两个以上需要重新计算辅助信息大小
    '''
    k = 16
    scale = k - 1
    for i in range(k - 2, -1, -1):
        if bins[i] == 250:
            scale = scale - 1
    print(scale)
    s2 = np.ceil(np.log2(T[scale - 1]))
    print('s2,scale', s2, scale)
    anx = cal_aux2(s2, scale, lcm, T, data_length)
    print('anx2', anx)
    ec1 = ecfinal / 2 + anx
    # # img_hide记录修改过后的图像坐标
    img_hide1 = copy.deepcopy(danci_img)
    # # 嵌入的数据secrt_msg[0][index]
    secret_info = secret_msg(row, col)
    pd.DataFrame(secret_info).to_csv('./search.csv')
    # # 嵌入数据之后的误差图像
    img_hide1, ec1 = insert_msg(em, Dbis, bins, img_hide1, secret_info, ec1, flags)
    # # lsb插入辅助信息
    anx_secret = secret_msg(row, col)
    img_hide1 = Lsb_insert(row, col, anx,img_hide1, flag, anx_secret)
    return img_hide1, ec1

def  plt_make(psnr):
    x = np.linspace(1, 6, 21)
    # x = np.arange(7)
    y = psnr
    plt.plot(x, y, '-r', label='psnr')
    plt.plot(x, y, '.b')
    plt.title('lena psnr')
    plt.xlabel('ec', color='#1C2833')
    plt.ylabel('psnr', color='#1C2833')
    plt.legend(loc='upper left')
    plt.grid()
    plt.show()
if __name__ == '__main__':

    ecfinal = 1e4
    x_size = np.linspace(1, 6, 21)
    y_size =np.zeros(21)
    index = 0
    for i in x_size:
        print("第{}论".format(i))
        ecfinal = i*1e4
        flag = 0
        k = 16
        Img = Load_img('test', flag)
        img = Img[0]
        imgs = copy.deepcopy(img)
        img_hide, ec = danci(img, flag, k, ecfinal)
        print("第{}论第一次插入结束".format(i))
        pd.DataFrame(img_hide).to_csv('./imgs.csv')
        new_img = copy.deepcopy(img_hide)
        print('ec',ec)
        if int(ec) == 0:
            k = 16
            flag = 1
            new_img, ec = danci(new_img, flag, k,ecfinal)
            pd.DataFrame(new_img).to_csv('./final.csv')
            pd.DataFrame(imgs).to_csv('./first.csv')
            psnr = cal_psnr(new_img, imgs,ec)
        print('psnr',psnr)
        y_size[index] = psnr
        index+=1
    pd.DataFrame(y_size).to_csv('./psnr.csv')
    plt_make(y_size)

    #
    # print(psnr)
    # print(em.shape)
    # # # 制作直方图之前需要降维
    # # print(em.flatten().count(-1))
    # index = 0
    # for i in em_update.flatten():
    #     if i == -2:
    #         index+=1
    # print(index)
    # plt_make(em.flatten(), '误差大小', '数量', '误差直方图')
    # plt_make(em_update.flatten(), '误差大小', '数量', '误差直方图2')
